# # 
#  crypt.py
#  ExpectIT
#  
#  Author: Marvin S. Addison
#  Date: 2007-10-15
#  Version: $Revision: $
#  Updated: $Date: $
#  Copyright 2008 Marvin S. Addison. All rights reserved.
#  See LICENSE.txt for licensing information.
#

import sys
import base64
from util import read_file, write_file
import Crypto.Cipher
from Crypto.Cipher import AES
from Crypto.Util.randpool import RandomPool

"""Initialization vector"""
IV = 'f0015901d155011d'

"""Default key length"""
DEFAULT_KEY_LENGTH = 2048

"""Entropy generation input message"""
ENTROPY_INPUT_MESSAGE = \
"""Type some random keys on the keyboard to generate entropy.
The more the merrier.  If you have children, this is a great time to
employ their natural entropy generation abilities.\n\n"""

def encrypt(keypath, plaintext):
	"""Encrypts a plaintext string with the symmetric key at the given path
	using the AES cipher.  The encrypted result is returned as a base-64
	encoded string."""
	ciphertext = get_cipher(keypath).encrypt(pad(plaintext, 16))
	return base64.encodestring(ciphertext)

def decrypt(keypath, ciphertext):
	"""Decrypts a base64-encoded ciphertext string with the 
	symmetric key at the given path using the AES cipher.  The resulting
	plaintext is returned."""
	bin_ciphertext = base64.decodestring(ciphertext)
	return get_cipher(keypath).decrypt(bin_ciphertext).rstrip('\0')

def genkey(keypath, keysize):
	"""Generates an AES key of the given length at the specified path.  Key
	size, in bytes, must be 16, 24, or 32."""
	if keysize % 16 != 0:
		raise Exception('AES key must be either 16, 24, or 32 bytes long.')
	rnd = RandomPool(keysize)
	rnd.add_event(raw_input(ENTROPY_INPUT_MESSAGE))
	write_file(keypath, rnd.get_bytes(keysize))
	
def get_cipher(keypath):
	"""Returns an AES cipher object."""
	return AES.new(read_file(keypath), AES.MODE_CBC, IV)

def pad(s, padlen):
	"""Pads a string with null characters so its length is a
	multiple of padlen."""
	while len(s) % padlen != 0:
		s += '\0'
	return s

def main(args):
	"""Main function that supports this module being called directly."""
	if len(args) < 3:
		print 'Not enough arguments.'
		usage()
	operation = args[0]
	try:
		method = getattr(sys.modules[__name__], operation)
	except AttributeError, ae:
		print 'Invalid operation "%s"' % operation
		usage()
	
	keypath = args[1]
	arg2 = args[2]
	if operation == 'genkey':
		try:
			arg2 = int(args[2])
		except ValueError:
			print '%s is not a valid integer' % arg2
			usage()
	try:
		retval = method(keypath, arg2)
		if retval:
			print '%s result: %s' % (operation, retval)
		print '%s operation completed.' % operation
	except IOError, ioe:
		print 'IO error.  Details:\n%s' % ioe
	except Exception, e:
		print 'Unexpected error. Details:\n%s' % e

def usage():
	print ''
	print 'USAGE: crypt operation arg1 arg2'
	print 'Supported operations:'
	print '  genkey keypath keysize'
	print '  encrypt keypath plaintext'
	print '  decrypt keypath ciphertext'
	sys.exit()

if __name__ == '__main__':
	if len(sys.argv) < 2:
		usage()
		sys.exit()
	if sys.argv[1] == 'crypt':
		main(sys.argv[2:])
	else:
		main(sys.argv[1:])
